import torch
import torchvision
import cv2
import numpy as np
import torch.nn.functional as F

class ScrewClassifier:
    def __init__(self,model_path=r'.\best_model.pth'):
        # 模型初始化
        self.model = torchvision.models.alexnet(pretrained=False)
        self.model.classifier[6] = torch.nn.Linear(self.model.classifier[6].in_features, 3)
        # 加载模型参数
        self.model.load_state_dict(torch.load(model_path))
        # 迁移到设备上
        self.device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self.model.to(self.device)
        # 正则化
        mean = 255.0 * np.array([0.48]) # 均值
        stdev = 255.0 * np.array([0.229, 0.224, 0.225]) # 基于样本估算的标准偏差
        self.normalize = torchvision.transforms.Normalize(mean, stdev)

    def __preprocess(self,input_img):
        # 预处理函数，用于转换图像通道等操作
        img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)
        img = img.transpose((2, 0, 1))
        img = torch.from_numpy(img).float()
        img = self.normalize(img)
        img = img.to(self.device)
        img = img[None, ...]
        return img

    def classify(self,img):
        x=self.__preprocess(img)
        result = self.model(x)
        result = F.softmax(result,dim=1)
        result = result.cpu().detach().numpy()
        return np.argmax(result)



if __name__ == '__main__':
    sc = ScrewClassifier()
    img=cv2.imread(r".\data\wrong\250.jpg")

    print(sc.classify(img))

